from octis.models.ETM import ETM
from octis.dataset.dataset import Dataset
from octis.evaluation_metrics.diversity_metrics import TopicDiversity
from octis.evaluation_metrics.coherence_metrics import Coherence

# dataset_name = "20NewsGroup"
dataset_name = "wos_mat"

# Define dataset
dataset = Dataset()
dataset.load_custom_dataset_from_folder("../preprocessed_datasets/" + dataset_name)

if __name__ == "__main__":
    model = ETM(num_topics=10, num_epochs=50)
    output = model.train_model(dataset)
    for t in output['topics'][:5]:
        print(" ".join(t))

    npmi = Coherence(texts=dataset.get_corpus(), topk=10, measure='c_npmi')
    Cv = Coherence(texts=dataset.get_corpus(), topk=10, measure='c_v')
    topic_diversity = TopicDiversity(topk=10)
    topic_diversity_score = topic_diversity.score(output)
    print("TD: " + str(topic_diversity_score))
    print("npmi: " + str(npmi.score(output)))
    print("Cv: " + str(Cv.score(output)))
